#!/usr/bin/env python3
"""
Register API 客户端演示
基于 motor/controller/api_server/controller_api.py 的 register 接口
"""

import requests
import json
import time
from typing import List, Dict, Any
from motor.resources.http_msg_spec import RegisterMsg, ServerInfo, Ranktable
from motor.resources.instance import ParallelConfig
from motor.resources.endpoint import DeviceInfo


class RegisterClient:
    """Register API 客户端类"""
    
    def __init__(self, controller_host: str = "localhost", controller_port: int = 8000, use_https: bool = False):
        """
        初始化客户端
        
        Args:
            controller_host: 控制器主机地址
            controller_port: 控制器端口
            use_https: 是否使用 HTTPS
        """
        self.protocol = "https" if use_https else "http"
        self.base_url = f"{self.protocol}://{controller_host}:{controller_port}"
        self.register_url = f"{self.base_url}/controller/register"
        self.heartbeat_url = f"{self.base_url}/controller/heartbeat"
        self.reregister_url = f"{self.base_url}/controller/reregister"
        
    def create_sample_register_msg(self, job_name: str = "test_job", role: str = "worker") -> RegisterMsg:
        """
        创建示例注册消息
        
        Args:
            job_name: 作业名称
            role: 角色 (worker, coordinator, etc.)
            
        Returns:
            RegisterMsg: 注册消息对象
        """
        # 创建设备信息
        device_info = DeviceInfo(
            device_id="0",
            device_ip="192.168.1.100",
            super_device_id=None,
            rank_id="0"
        )
        
        # 创建服务器信息
        server_info = ServerInfo(
            server_id="server_0",
            host_ip="192.168.1.100",
            device=[device_info]
        )
        
        # 创建排名表
        ranktable = Ranktable(
            version="1.0",
            status="active",
            server_count="1",
            server_list=[server_info]
        )
        
        # 创建并行配置 - 现在可以使用字段名或简短参数名
        parallel_config = ParallelConfig(
            dp_size=1,  # 数据并行大小
            cp_size=1,  # 上下文并行大小
            tp_size=1,  # 张量并行大小
            sp_size=1,  # 序列并行大小
            ep_size=1,  # 专家并行大小
            pp_size=1,  # 流水线并行大小
            world_size=1
        )
        
        # 创建注册消息
        register_msg = RegisterMsg(
            job_name=job_name,
            model_name="test_model",
            role=role,
            pod_ip="127.0.0.1",
            host_ip="192.168.1.100",
            bussiness_port=["8080"],
            mgmt_port="9090",
            parallel_config=parallel_config,
            ranktable=ranktable
        )
        
        return register_msg
    
    def register(self, register_msg: RegisterMsg) -> Dict[str, Any]:
        """
        发送注册请求
        
        Args:
            register_msg: 注册消息对象
            
        Returns:
            Dict: 服务器响应
        """
        try:
            # 将 Pydantic 模型转换为字典
            data = register_msg.model_dump()
            
            print(f"发送注册请求到: {self.register_url}")
            print(f"注册数据: {json.dumps(data, indent=2, ensure_ascii=False)}")
            
            # 发送 POST 请求
            response = requests.post(
                self.register_url,
                json=data,
                timeout=30,
                verify=False if self.protocol == "https" else True
            )
            
            # 检查响应状态
            response.raise_for_status()
            
            result = response.json()
            print(f"注册成功! 响应: {json.dumps(result, indent=2, ensure_ascii=False)}")
            return result
            
        except requests.exceptions.RequestException as e:
            print(f"注册请求失败: {e}")
            if hasattr(e, 'response') and e.response is not None:
                print(f"错误响应: {e.response.text}")
            raise
        except Exception as e:
            print(f"注册过程中发生错误: {e}")
            raise
    
    def register_with_custom_data(self, 
                                 job_name: str,
                                 model_name: str,
                                 role: str,
                                 pod_ip: str,
                                 host_ip: str,
                                 business_ports: List[str],
                                 mgmt_port: str,
                                 device_infos: List[Dict[str, Any]] = None) -> Dict[str, Any]:
        """
        使用自定义数据注册
        
        Args:
            job_name: 作业名称
            role: 角色
            pod_ip: Pod IP
            host_ip: 主机 IP
            business_ports: 业务端口列表
            mgmt_port: 管理端口
            device_infos: 设备信息列表
            
        Returns:
            Dict: 服务器响应
        """
        # 默认设备信息
        if device_infos is None:
            device_infos = [
                {
                    "device_id": "0",
                    "device_ip": host_ip,
                    "super_device_id": None,
                    "rank_id": "0"
                }
            ]
        
        # 创建设备信息对象
        devices = [DeviceInfo(**info) for info in device_infos]
        
        # 创建服务器信息
        server_info = ServerInfo(
            server_id=f"server_{job_name}",
            host_ip=host_ip,
            device=devices
        )
        
        # 创建排名表
        ranktable = Ranktable(
            version="1.0",
            status="active",
            server_count="1",
            server_list=[server_info]
        )
        
        # 创建并行配置 - 现在可以使用字段名
        parallel_config = ParallelConfig(
            dp_size=1,
            cp_size=1,
            tp_size=1,
            sp_size=1,
            ep_size=1,
            pp_size=1,
            world_size=len(device_infos)  # 根据设备数量设置 world_size
        )
        
        # 创建注册消息
        register_msg = RegisterMsg(
            job_name=job_name,
            model_name=model_name,
            role=role,
            pod_ip=pod_ip,
            host_ip=host_ip,
            bussiness_port=business_ports,
            mgmt_port=mgmt_port,
            parallel_config=parallel_config,
            ranktable=ranktable
        )
        
        return self.register(register_msg)


def main():
    """主函数 - 演示如何使用 RegisterClient"""
    
    print("=== Register API 客户端演示 ===\n")
    
    # 创建客户端实例
    client = RegisterClient(
        controller_host="localhost",
        controller_port=8000,
        use_https=False  # 设置为 True 如果使用 HTTPS
    )
    
    try:
        # 方法1: 使用示例数据注册
        print("1. 使用示例数据注册:")
        sample_msg = client.create_sample_register_msg("demo_job", "worker")
        result1 = client.register(sample_msg)
        print(f"注册结果: {result1}\n")
        
        # 方法2: 使用自定义数据注册
        print("2. 使用自定义数据注册:")
        result2 = client.register_with_custom_data(
            job_name="custom_job",
            model_name="custom_model",
            role="coordinator",
            pod_ip="127.0.0.1",
            host_ip="192.168.1.101",
            business_ports=["8080"],
            mgmt_port="9091",
            device_infos=[
                {
                    "device_id": "0",
                    "device_ip": "192.168.1.101",
                    "super_device_id": None,
                    "rank_id": "0"
                },
                {
                    "device_id": "1", 
                    "device_ip": "192.168.1.101",
                    "super_device_id": None,
                    "rank_id": "1"
                }
            ]
        )
        print(f"自定义注册结果: {result2}\n")
        
        # 方法3: 批量注册多个实例
        print("3. 批量注册多个实例:")
        jobs = [
            {"job_name": "job_1", "role": "worker", "pod_ip": "127.0.0.1"},
            {"job_name": "job_2", "role": "worker", "pod_ip": "127.0.0.1"},
            {"job_name": "job_3", "role": "coordinator", "pod_ip": "127.0.0.1"}
        ]
        
        for i, job in enumerate(jobs):
            print(f"注册作业 {i+1}: {job['job_name']}")
            result = client.register_with_custom_data(
                job_name=job["job_name"],
                model_name="model_1",
                role=job["role"],
                pod_ip=job["pod_ip"],
                host_ip=f"192.168.1.{102+i}",
                business_ports=[f"808{i}"],
                mgmt_port=f"909{i}"
            )
            print(f"作业 {job['job_name']} 注册结果: {result}")
            time.sleep(1)  # 避免请求过快
            
    except Exception as e:
        print(f"演示过程中发生错误: {e}")
        print("请确保控制器服务正在运行并且可以访问")


if __name__ == "__main__":
    main()
